from metagpt.schema import Message
from metagpt.roles import Role
from asyncio import Queue, gather, create_task
from typing import Dict
import asyncio

# 全局消息路由
class MessageRouter:
    _queues: Dict[str, Queue] = {}
    
    @classmethod
    def get_queue(cls, name: str) -> Queue:
        name = cls._ensure_string(name)
        if name not in cls._queues:
            cls._queues[name] = Queue()
        return cls._queues[name]
    
    @classmethod
    def _ensure_string(cls, name) -> str:
        """确保名称是字符串"""
        if isinstance(name, (set, list, tuple)):
            return next(iter(name))
        return str(name)
    
    @classmethod
    async def publish(cls, message: Message):
        # 强制转换 send_to 为字符串
        send_to = cls._ensure_string(message.send_to)
        sent_from = cls._ensure_string(message.sent_from)
        
        print(f"[Router] Routing message from {sent_from} to {send_to}")
        await cls.get_queue(send_to).put(message)

class TaskManager(Role):
    def __init__(self, name="TaskManager"):
        super().__init__(name=name)
        self.router = MessageRouter
    
    async def run(self):
        try:
            print(f"[{self.name}] Starting TaskManager...")
            
            # 创建消息时强制转换参数
            task_a = self._create_message("Process data A", "WorkerA")
            task_b = self._create_message("Process data B", "WorkerB")
            
            print(f"[{self.name}] Sending tasks to workers...")
            await self.router.publish(task_a)
            await self.router.publish(task_b)
            
            # 等待结果
            print(f"[{self.name}] Waiting for results...")
            results = await gather(
                self._wait_for_response("WorkerA"),
                self._wait_for_response("WorkerB"),
                return_exceptions=True
            )
            
            self._process_results(results)
                
        except Exception as e:
            print(f"[{self.name}] Critical error: {e}")

    def _create_message(self, content: str, send_to: str) -> Message:
        """创建消息并确保参数类型正确"""
        return Message(
            content=content,
            role=self.name,
            send_to=str(send_to),  # 强制转换为字符串
            sent_from=str(self.name)
        )

    async def _wait_for_response(self, worker: str) -> Message:
        print(f"[{self.name}] Waiting for response from {worker}...")
        while True:
            msg = await self.router.get_queue(self.name).get()
            if str(msg.sent_from) == worker:  # 强制类型比较
                print(f"[{self.name}] Received response from {worker}")
                return msg
            else:
                create_task(self.router.get_queue(self.name).put(msg))
                await asyncio.sleep(0.1)

    def _process_results(self, results):
        print(f"[{self.name}] Results received:")
        for worker, result in zip(["WorkerA", "WorkerB"], results):
            if isinstance(result, Exception):
                print(f"  Error from {worker}: {str(result)}")
            elif result:
                print(f"  From {worker}: {result.content}")
            else:
                print(f"  No response from {worker}")

class Worker(Role):
    def __init__(self, name: str):
        super().__init__(name=name)
        self.router = MessageRouter
    
    async def run(self):
        try:
            print(f"[{self.name}] Starting...")
            
            # 等待任务
            print(f"[{self.name}] Waiting for task...")
            task = await self.router.get_queue(self.name).get()
            print(f"[{self.name}] Received task: {task.content}")
            
            # 处理并回复
            await self._process_task(task)
            
        except Exception as e:
            print(f"[{self.name}] Error: {e}")

    async def _process_task(self, task: Message):
        result = f"Processed {task.content}"
        response = Message(
            content=result,
            role=self.name,
            send_to=str(task.sent_from),  # 强制转换
            sent_from=str(self.name)
        )
        print(f"[{self.name}] Sending result to {response.send_to}")
        await self.router.publish(response)

async def main():
    print("Initializing system...")
    
    # 创建角色实例
    roles = [
        TaskManager(),
        Worker("WorkerA"),
        Worker("WorkerB")
    ]
    
    # 运行所有角色
    await gather(*[role.run() for role in roles], return_exceptions=True)

if __name__ == "__main__":
    print("Starting program...")
    asyncio.run(main())
